from langchain_redis import RedisChatMessageHistory
from langchain_core.messages import BaseMessage, AIMessage, ToolMessage
from datetime import datetime
from redis.commands.json.path import Path


class ZLRedisChatMessageHistory(RedisChatMessageHistory):
    def add_message(self, message: BaseMessage) -> None:
        data = {
            "content": message.content,
            "additional_kwargs": message.additional_kwargs,
            "type": message.type,
        }
        if isinstance(message, ToolMessage):
            data['tool_call_id'] = message.tool_call_id
        elif isinstance(message, AIMessage) and message.tool_calls:
            data['tool_calls'] = message.tool_calls

        data_to_store = {
            "type": message.type,
            "data": data,
            "session_id": self.session_id,
            "timestamp": datetime.now().timestamp(),
        }

        key = f"{self.key_prefix}{self.session_id}:{data_to_store['timestamp']}"
        self.redis_client.json().set(key, Path.root_path(), data_to_store)

        if self.ttl:
            self.redis_client.expire(key, self.ttl)
